# 导入聚宽核心库（策略运行所需的API）
from jqdata import *

# ========== 1) 初始化 ==========
def initialize(context):
    set_benchmark('000300.XSHG')                                   # 设定沪深300为业绩基准
    set_option('use_real_price', True)                              # 交易与回测均使用实时价格/当前价
    log.set_level('order', 'error')                                 # 减少下单日志噪音，保留报错

    # 设定滑点与佣金（相对保守，贴近实际）
    set_slippage(PriceRelatedSlippage(0.002))                       # 按价格比例滑点：0.2%
    set_order_cost(OrderCost(open_tax=0, close_tax=0.001,           # A股卖出印花税千一
                             open_commission=0.0003, close_commission=0.0003,
                             min_commission=5), type='stock')       # 佣金与最小费用

    # 全局参数（g 对象）
    g.index = '000300.XSHG'                                         # 择时参考指数：沪深300
    g.max_hold = 10                                                 # 最多持股数量：10只，便于分散又不至于过度分散
    g.ipo_days = 60                                                 # 次新股剔除天数：60日
    g.min_price = 3                                                 # 低价股过滤门槛：3元
    g.stop_loss_pct = -0.08                                         # 单票止损：-8%
    g.rebalance_time = '09:45'                                      # 调仓时间：开盘后避开集合竞价
    g.stop_time = '14:50'                                           # 止损检查时间：尾盘前

    # 定时任务：每日例行
    run_daily(rebalance, time=g.rebalance_time)                     # 每日9:45 调仓（含择时与选股）
    run_daily(stop_loss_check, time=g.stop_time)                    # 每日14:50 止损检查


# ========== 2) 调仓主流程 ==========
def rebalance(context):
    if not market_timing_ok(context):                               # 简单双均线择时（指数多头才开仓）
        for s in list(context.portfolio.positions.keys()):          # 若不满足择时，清仓防守
            order_target_percent(s, 0)
        return                                                      # 退出，不做新开仓

    universe = build_universe(context)                              # 构建可交易股票池（流动性+交易约束）
    if len(universe) == 0:                                          # 防御性处理
        return

    target_list, weights = select_and_rank(context, universe)       # 多因子打分 + 动量排序，得到买入清单与权重
    if len(target_list) == 0:
        return

    # 先卖后买，释放资金
    current = set(context.portfolio.positions.keys())               # 当前持仓集合
    target = set(target_list)                                       # 目标持仓集合
    for s in current - target:                                      # 卖出不在目标中的持仓
        if is_tradeable(s):
            order_target_percent(s, 0)

    # 等权买入目标标的
    for s in target_list:
        if is_tradeable(s):
            order_target_percent(s, weights.get(s, 0))


# ========== 3) 止损与破位检查 ==========
def stop_loss_check(context):
    for s, pos in context.portfolio.positions.items():              # 遍历持仓
        if pos.avg_cost <= 0:                                       # 无有效成本价（新成交等），跳过
            continue
        cd = get_current_data()
        if s not in cd or cd[s].paused:                             # 停牌则跳过
            continue
        price = cd[s].last_price                                    # 当前价
        if price is None:
            continue

        # 1) 固定比例止损：回撤超过8%则清仓
        if (price / pos.avg_cost - 1) <= g.stop_loss_pct:
            if is_tradeable(s):
                order_target_percent(s, 0)
            continue

        # 2) 技术破位止损：跌破近10日最低价则清仓
        lows = get_bars(s, count=11, unit='1d', fields=['low'])     # 取近11日的最低价序列
        if lows is not None and len(lows) >= 11:
            recent_low = float(min(lows['low'][:-1]))               # 最近10日（不含当日）最低
            if price < recent_low and is_tradeable(s):
                order_target_percent(s, 0)


# ========== 4) 择时：指数收盘在MA60上方 ==========
def market_timing_ok(context):
    # 使用指数近65日收盘，计算MA60
    df = get_price(g.index, count=65, end_date=context.previous_date, fields=['close'])
    if df is None or len(df) < 60:                                  # 数据不足时不做多
        return False
    ma60 = df['close'].rolling(60).mean().iloc[-1]                  # 60日均线
    last = df['close'].iloc[-1]                                     # 最新收盘
    return last > ma60                                              # 收盘>MA60视为多头


# ========== 5) 股票池构建：A股交易约束与流动性 ==========
def build_universe(context):
    # 初级流动性：沪深300 + 中证500（覆盖大中盘，增强可成交性）
    pool = set(get_index_stocks('000300.XSHG')) | set(get_index_stocks('000905.XSHG'))

    # 基础过滤：去ST、停牌、次新、低价、涨跌停
    cd = get_current_data()
    today = context.current_dt.date()
    res = []
    for s in pool:
        # 证券基础信息（上市时间）
        info = get_security_info(s)
        if info is None:
            continue
        if (today - info.start_date).days < g.ipo_days:             # 次新剔除
            continue
        if s not in cd:
            continue
        if cd[s].paused or cd[s].is_st:                             # 停牌或ST剔除
            continue
        lp = cd[s].last_price
        if lp is None or lp < g.min_price:                          # 低价股剔除（含无价）
            continue
        # 涨跌停附近不下单（提高成交确定性）
        if lp >= cd[s].high_limit * 0.999 or lp <= cd[s].low_limit * 1.001:
            continue
        res.append(s)
    return res


# ========== 6) 多因子选股 + 动量排序 ==========
def select_and_rank(context, stocks):
    if len(stocks) == 0:
        return [], {}

    # 取基本面因子：ROE、毛利率、营收同比；估值：PE、PB
    q = query(valuation.code,
              valuation.pe_ratio, valuation.pb_ratio,
              indicator.roe, indicator.gross_profit_margin,
              indicator.inc_revenue_year_on_year
             ).filter(valuation.code.in_(stocks))
    df = get_fundamentals(q)                                        # 返回DataFrame，index为数字，含code列
    if df is None or df.empty:
        return [], {}

    df = df.set_index('code')                                       # 将股票代码作为索引
    # 缺失值填充与合理裁剪
    df['pe_ratio'] = df['pe_ratio'].clip(lower=0).replace([float('inf')], None)
    df['pb_ratio'] = df['pb_ratio'].clip(lower=0).replace([float('inf')], None)
    df['roe'] = df['roe'].fillna(0)
    df['gross_profit_margin'] = df['gross_profit_margin'].fillna(0)
    df['inc_revenue_year_on_year'] = df['inc_revenue_year_on_year'].fillna(0)
    df = df.dropna(subset=['pe_ratio', 'pb_ratio'])                 # 估值必须可用

    # 动量：近60日涨跌幅（剔除当日，避免信号-成交冲突）
    moms = {}
    for s in df.index:
        p = get_price(s, count=61, end_date=context.previous_date, fields=['close'])
        if p is None or len(p) < 2:
            moms[s] = 0
        else:
            moms[s] = float(p['close'].iloc[-1] / p['close'].iloc[0] - 1)
    df['mom60'] = pd.Series(moms)

    # 计算z-score（标准化）并构建综合得分
    def zscore(series):
        m, st = series.mean(), series.std()
        if st is None or st == 0 or pd.isna(st):
            return series * 0                                       # 全零，避免除零
        return (series - m) / st

    df['score'] = (
        zscore(df['roe']) +
        zscore(df['gross_profit_margin']) +
        zscore(df['inc_revenue_year_on_year']) -
        zscore(df['pe_ratio']) -
        zscore(df['pb_ratio']) +
        zscore(df['mom60'])
    )

    # 进一步以流动性偏好（大市值）筛前300，再从中选Top N，提升成交确定性
    q2 = query(valuation.code, valuation.circulating_market_cap
              ).filter(valuation.code.in_(list(df.index)))
    cap_df = get_fundamentals(q2).set_index('code')
    df = df.join(cap_df, how='inner')
    df = df.sort_values(by='circulating_market_cap', ascending=False).head(300)

    # 最终按score排序，取前g.max_hold只，等权配置
    df = df.sort_values(by='score', ascending=False)
    picks = list(df.index[:g.max_hold])
    if len(picks) == 0:
        return [], {}

    w = 1.0 / len(picks)
    weights = {s: w for s in picks}
    return picks, weights


# ========== 7) 交易可行性检查（停牌/涨跌停） ==========
def is_tradeable(s):
    cd = get_current_data()
    if s not in cd:
        return False
    if cd[s].paused:
        return False
    lp = cd[s].last_price
    if lp is None:
        return False
    # 避免涨停封板或接近涨停/跌停价的不可成交风险
    if lp >= cd[s].high_limit * 0.999:
        return False
    if lp <= cd[s].low_limit * 1.001:
        return False
    return True
